import os
import shutil
from MDtools.lmp import changecode

import os
import shutil
import random

def mission_creat(mission_space_name, mission_condi, 
                 change_file='input.lammps', random_num=None):
    """
    创建一个任务空间，根据指定的参数条件和可选的随机子目录。

    参数:
    - mission_space_name (str): 要创建的任务空间目录名称。
    - mission_condi (dict): 参数及其可能取值的字典。
    - change_file (str): 要修改的文件名（默认: 'input.lammps'）。
    - random_num (int, optional): 在每个最低层级目录下创建的随机子目录数量。

    返回:
    - list: 所有创建的任务目录路径列表。
    """

    
    # 1. 获取当前目录和源文件夹路径
    current_dir = os.getcwd()
    source_dir = os.path.join(current_dir, 'source_file')

    if not os.path.exists(source_dir):
        raise FileNotFoundError(f"源目录 '{source_dir}' 不存在。")
    
    # 2. 创建任务空间目录在父目录下
    parent_dir = os.path.dirname(current_dir)
    mission_space_path = os.path.join(parent_dir, mission_space_name)
    os.makedirs(mission_space_path, exist_ok=True)
    
    # 获取参数及其取值列表
    params = list(mission_condi.keys())
    param_values_list = [mission_condi[param] for param in params]
    
    # 递归函数：创建目录并复制文件
    def create_dirs(current_path, params_remaining, param_values_remaining, param_values_combination):
        
        clear_local_directory(current_path)
        
        if not params_remaining:
            # 到达最低层级
            if random_num and isinstance(random_num, int):
                # 1. 生成 random_num 个6位随机整数作为种子
                seeds = [random.randint(100000, 999999) for _ in range(random_num)]
                
                randomized_paths = []
                for i, seed in enumerate(seeds):
                    sub_dir_name = f'N{i}'
                    sub_dir_path = os.path.join(current_path, sub_dir_name)
                    os.makedirs(sub_dir_path, exist_ok=True)
                    
                    # 2. 直接从源文件夹复制所有文件到 N* 文件夹
                    for item in os.listdir(source_dir):
                        s = os.path.join(source_dir, item)
                        d = os.path.join(sub_dir_path, item)
                        if os.path.isfile(s):
                            shutil.copy2(s, d)
                        elif os.path.isdir(s):
                            shutil.copytree(s, d, dirs_exist_ok=True)
                    
                    # 3. 修改 N* 文件夹中的 input.lammps 文件的 'seed' 参数
                    infile_dest = os.path.join(sub_dir_path, change_file)
                    if os.path.exists(infile_dest):
                        # 首先根据参数组合修改其他参数
                        for param, value in param_values_combination.items():
                            changecode(infile_dest, param, value)
                        # 然后修改 'seed' 参数
                        changecode(infile_dest, 'seed', seed)
                    else:
                        raise FileNotFoundError(f"'{change_file}' 在 '{sub_dir_path}' 中不存在。")
                    
                    randomized_paths.append(sub_dir_path)
                
                return randomized_paths
            else:
                # random_num 未指定，直接复制文件到 current_path 并修改参数
                for item in os.listdir(source_dir):
                    s = os.path.join(source_dir, item)
                    d = os.path.join(current_path, item)
                    if os.path.isfile(s):
                        shutil.copy2(s, d)
                    elif os.path.isdir(s):
                        shutil.copytree(s, d, dirs_exist_ok=True)
                
                # 修改 input.lammps 文件中的参数
                infile_dest = os.path.join(current_path, change_file)
                if os.path.exists(infile_dest):
                    for param, value in param_values_combination.items():
                        changecode(infile_dest, param, value)
                else:
                    raise FileNotFoundError(f"'{change_file}' 在 '{current_path}' 中不存在。")
                
                return [current_path]
        else:
            # 还未到最低层级，继续创建子目录
            param = params_remaining[0]
            values = param_values_remaining[0]
            all_paths = []
            for value in values:
                if type(value) == type('1') and ' ' in value: # 防止路径里面有空格，这个是一个愚昧的历史问题。
                    value=value.replace(' ','')
                dir_name = f'{param}={value}'
                dir_path = os.path.join(current_path, dir_name)
                os.makedirs(dir_path, exist_ok=True)
                
                # 更新参数组合
                new_param_values_combination = param_values_combination.copy()
                new_param_values_combination[param] = value
                
                # 递归调用
                paths = create_dirs(
                    dir_path,
                    params_remaining[1:],
                    param_values_remaining[1:],
                    new_param_values_combination
                )
                all_paths.extend(paths)
            return all_paths
    
    # 开始递归创建目录
    created_mission_paths = create_dirs(mission_space_path, params, param_values_list, {})
    
    return created_mission_paths

def mission_sub(mission_space_path):

    import os
    import subprocess

    # Save the original working directory
    original_dir = os.getcwd()
    for dir in mission_space_path:
        os.chdir(dir)
        subprocess.run(['bash', 'run'])
        # Return to the original directory
        os.chdir(original_dir)

def mission_analy(mission_space_path, analysis_file):

    # 在 analysis_file 的上上级路径输出
    root_dir = os.path.dirname(os.path.dirname(analysis_file))
    result_dir = os.path.join(root_dir, 'result')

    # 创建 result 目录（如果不存在）
    os.makedirs(result_dir, exist_ok=True)

    # 如果 result_dir 存在且不是空文件夹，删除其中的所有内容
    if os.path.exists(result_dir) and os.path.isdir(result_dir):
        for item in os.listdir(result_dir):
            item_path = os.path.join(result_dir, item)
            # 如果是文件，删除文件；如果是目录，递归删除目录
            if os.path.isfile(item_path):
                os.remove(item_path)
            elif os.path.isdir(item_path):
                shutil.rmtree(item_path)

    analysis_filename = os.path.basename(analysis_file)
    target_dirs = []
    target_miss_nums = []
    # Step 1: Traverse directories and stop at the level containing change_file
    for root, dirs, files in os.walk(mission_space_path):
        if 'input.lammps' in files:
            target_dirs.append(root)
            target_miss_nums.append(find_mission_number(files))
            # Remove the 'analysis' folder if it exists
            analysis_dir = os.path.join(root, 'analysis')
            if os.path.exists(analysis_dir):
                shutil.rmtree(analysis_dir)
                #print(f"Removed {analysis_dir}")
            # Remove the 'his' folder if it exists
            his_dir = os.path.join(root, 'his')
            if os.path.exists(his_dir):
                shutil.rmtree(his_dir)
                #print(f"Removed {his_dir}")
            
            # Copy the analysis_file into this directory
            if analysis_filename in files:
                shutil.copy2(analysis_file, root)
                #print(f"Copied {analysis_file} to {root}")

    # Step 2: Execute 'bash _analysis.sh' in each target directory
    for target_dir,target_miss_num in zip(target_dirs,target_miss_nums):
        os.chdir(target_dir)
        os.system('bash _analysis.sh')
    
    # Collect and rename new 'analysis' files
        analysis_file_path = os.path.join(target_dir, 'analysis')
        if os.path.exists(analysis_file_path):
            # Create a new filename by replacing '/' with '_'
            relative_path = os.path.relpath(target_dir, mission_space_path)
            new_filename = target_miss_num +'_'+ relative_path.replace(os.sep, '_') 
            new_filepath = os.path.join(result_dir, new_filename)
            # Move and rename the file
            shutil.copytree(analysis_file_path, new_filepath, dirs_exist_ok=True)
            print(f"{new_filepath} is done. ")


def find_mission_number(file_list):
    import re
    # 2024-11-29
    # 使用正则表达式提取所有符合 "_out.lammps" 格式的文件名中的数字部分
    numbers = []
    for filename in file_list:
        match = re.match(r'(\d+)_out\.lammps', filename)
        if match:
            numbers.append(match.group(1))  # 提取数字并转换为整数
    
    # 返回最大数字
    if numbers:
        return max(numbers)
    else:
        return None  # 如果没有找到匹配的文件名
    


def clear_local_directory(local_directory):
    """
    Prompts the user for confirmation, then clears all files and subdirectories
    in the specified local directory if confirmed.

    :param local_directory: The path of the local directory to be cleared.
    """
    if not os.path.exists(local_directory):
        print(f"Directory does not exist: {local_directory}")
        return

    if not os.listdir(local_directory):
        return

    # Prompt the user for confirmation
    user_input = input(f"The directory '{local_directory}' is not empty. Do you want to clear it? (yes/no): ").strip().lower()
    if user_input not in ['y','yes']:
        return

    # Proceed to clear the directory
    os.system(f'rm -rf  {local_directory}/* ')